use std::{net::IpAddr, sync::Arc};

use axum::{
    async_trait,
    extract::{FromRequest, FromRequestParts, Request},
    http::{request::Parts, StatusCode},
};
use axum_client_ip::{InsecureClientIp, SecureClientIp};

use crate::AppState;

/// 获得客户ip地址, 获取失败返回None
/// - 配置了ip_header, 则只从对应header获取，ip_header从Arc<AppState>中获取
/// - 未配置ip_header, 则从各类proxy_header依次获取直到ConnectInfo
/// - ConnectInfo 需要 启用[`axum::routing::Router::into_make_service_with_connect_info`]
pub struct ClientIp(pub Option<IpAddr>);

type Rejection = (StatusCode, &'static str);

#[async_trait]
impl<S> FromRequestParts<S> for ClientIp
where
    S: Send + Sync,
{
    type Rejection = Rejection;

    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
        let ip = if let Some(s) = parts.extensions.get::<Arc<AppState>>() {
            if let Some(ip) = s.ip_header() {
                SecureClientIp::from(ip, &parts.headers, &parts.extensions)
                    .map(|v| v.0)
                    .ok()
            } else {
                InsecureClientIp::from(&parts.headers, &parts.extensions)
                    .map(|v| v.0)
                    .ok()
            }
        } else {
            InsecureClientIp::from(&parts.headers, &parts.extensions)
                .map(|v| v.0)
                .ok()
        };
        Ok(ClientIp(ip))
    }
}
